import asyncio
import pandas as pd
import io
import datetime
import random
from typing import Optional

import requests
from fastmcp import FastMCP
from fastmcp import Client

mcp = FastMCP("demo.mcp")


@mcp.tool()
async def get_nasdaq_earnings_calendar(
    date: str | None = None,
    limit: int = 100
) -> str:
    """Get earnings calendar for a specific date using Nasdaq API.
    Date in YYYY-MM-DD format (defaults to today)
    Returns CSV with: Date, Symbol, Company Name, EPS, % Surprise, Market Cap, etc.
    Note: Single date only - call multiple times for date ranges.
    """
    # Constants
    BROWSER_HEADERS = {
    'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36'
    }
    NASDAQ_EARNINGS_URL = "https://api.nasdaq.com/api/calendar/earnings"
    NASDAQ_HEADERS = {
        **BROWSER_HEADERS,
        'Referer': 'https://www.nasdaq.com/'
    }
    

    # Set default date if not provided or validate provided date
    date_str = date if date else datetime.date.today().strftime('%Y-%m-%d')
    url = f"{NASDAQ_EARNINGS_URL}?date={date_str}"

    try:
        resp = requests.get(url, headers=NASDAQ_HEADERS, timeout=10)
        resp.raise_for_status()
        data = resp.json()

        if 'data' in data and data['data']:
            earnings_data = data['data']

            if earnings_data.get('headers') and earnings_data.get('rows'):
                column_headers = earnings_data['headers']
                rows = earnings_data['rows']

                # Extract column names from headers dict
                if isinstance(column_headers, dict):
                    column_names = list(column_headers.values())
                    column_keys = list(column_headers.keys())
                else:
                    column_names = [h.get('label', h) if isinstance(h, dict) else str(h) for h in column_headers]
                    column_keys = column_names

                # Convert rows to DataFrame
                processed_rows = []
                for row in rows:
                    if isinstance(row, dict):
                        processed_row = [row.get(key, '') for key in column_keys]
                        processed_rows.append(processed_row)

                if processed_rows:
                    df = pd.DataFrame(processed_rows, columns=column_names)
                    # Add date column at the beginning
                    df.insert(0, 'Date', date_str)

                    # Apply limit
                    if len(df) > limit:
                        df = df.head(limit)

                    output = io.StringIO()
                    df.to_csv(output, index=False)
                    return output.getvalue()

        # No earnings data found
        return f"No earnings announcements found for {date_str}. This could be due to weekends, holidays, or no scheduled earnings on this date."

    except Exception as e:
        return f"Error retrieving earnings data for {date_str}: {str(e)}"
    
GET_NASDAQ_EARNINGS_CALENDAR_SCHEMA = {
    "type": "function",
    "function": {
        "name": "get_nasdaq_earnings_calendar",
        "description": "Retrieve the Nasdaq earnings calendar for a specific date using the official Nasdaq API. Returns company ticker, name, EPS, surprise percentage, market cap, and more in CSV format.",
        "parameters": {
            "type": "object",
            "properties": {
                "date": {
                    "type": "string",
                    "description": "The date to query, in YYYY-MM-DD format. Defaults to today's date if not provided. Only single-day queries are supported."
                },
                "limit": {
                    "type": "integer",
                    "description": "The maximum number of results to return. Used to limit the number of CSV rows. Default is 100."
                }
            },
            "required": []
        }
    }
}

# --- 测试入口 ---
async def main():
    client = Client(mcp)
    async with client:
        tools = await client.list_tools()
        print("Available tools:", tools)

        result = await client.call_tool("get_nasdaq_earnings_calendar", {"date": "2025-10-15", "limit": 5})
        print("Earnings Calendar Result:\n", result)


if __name__ == "__main__":
    asyncio.run(main())